package apps.demo;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;

import be.ac.ulg.montefiore.run.jahmm.Hmm;
import be.ac.ulg.montefiore.run.jahmm.Observation;
import be.ac.ulg.montefiore.run.jahmm.ObservationDiscrete;
import be.ac.ulg.montefiore.run.jahmm.OpdfDiscrete;
import be.ac.ulg.montefiore.run.jahmm.OpdfDiscreteFactory;
import be.ac.ulg.montefiore.run.jahmm.draw.GenericHmmDrawerDot;
import be.ac.ulg.montefiore.run.jahmm.learn.BaumWelchLearner;
import be.ac.ulg.montefiore.run.jahmm.toolbox.KullbackLeiblerDistanceCalculator;
import be.ac.ulg.montefiore.run.jahmm.toolbox.MarkovGenerator;
import be.ac.ulg.montefiore.run.jahmm.ViterbiCalculator;

public class PresentationDemo {
	private enum Weather {
		Sunny,
		Cloudy,
		Rainy;
		
		private static Map<Weather, Integer> enumToIntMap = new HashMap<Weather, Integer>();
		private static Map<Integer, Weather> intToEnumMap = new HashMap<Integer, Weather>();
		private static Map<String, Weather> stringToEnumMap = new HashMap<String, Weather>();
		static {
			int i = 0;
			for (Weather e : values()) {
					enumToIntMap.put(e, i);
					intToEnumMap.put(i, e);
					
					stringToEnumMap.put(e.name().toUpperCase(), e);
					
					i++;
			}
		}
		
		public static Weather fromInt(int value) {
			return intToEnumMap.get(value);
		}

		public int toInt() {
			return enumToIntMap.get(this);
		}
		
		public long getLong() {
			return toInt();
		}
		
		public static Weather parseLong(long val) {
			return fromInt((int)val);
		}
		
		public static Weather fromString(String value) {
			return stringToEnumMap.get(value.toUpperCase());
		}
		
		public String toString() {
			return this.name();
		}
	};
	
	private enum Packet {
		Dry, 
		Dryish,
		Damp,
		Soggy
		;

		public ObservationDiscrete<Packet> observation() {
			return new ObservationDiscrete<Packet>(this);
		}
		
		private static Map<Packet, Integer> enumToIntMap = new HashMap<Packet, Integer>();
		private static Map<Integer, Packet> intToEnumMap = new HashMap<Integer, Packet>();
		private static Map<String, Packet> stringToEnumMap = new HashMap<String, Packet>();
		static {
			int i = 0;
			for (Packet e : values()) {
					enumToIntMap.put(e, i);
					intToEnumMap.put(i, e);
					
					stringToEnumMap.put(e.name().toUpperCase(), e);
					
					i++;
			}
		}
		
		public static Packet fromInt(int value) {
			return intToEnumMap.get(value);
		}

		public int toInt() {
			return enumToIntMap.get(this);
		}
		
		public long getLong() {
			return toInt();
		}
		
		public static Packet parseLong(long val) {
			return fromInt((int)val);
		}
		
		public static Packet fromString(String value) {
			return stringToEnumMap.get(value.toUpperCase());
		}
		
		public String toString() {
			return this.name();
		}
	};

	static {
		System.out.print("\n+++++++++++++++++++++++++++++++++++++++++++++++\n");
		System.out.print("\nBuilt In Model\n");
		System.out.print("" + getDefaultHmm());
		System.out.print("\n+++++++++++++++++++++++++++++++++++++++++++++++\n\n\n");
	}
	
	static Hmm<ObservationDiscrete<Packet>> defaultHmm = null;
	static Hmm<ObservationDiscrete<Packet>> getDefaultHmm() {
		if (defaultHmm == null) {
			defaultHmm = buildHmm();
		}
		return defaultHmm;
	}
	
	static List<ObservationDiscrete<Packet>> defaultSequences = null;
	static List<ObservationDiscrete<Packet>> getDefaultSequences() {
		if (defaultSequences == null) {
			ObservationDiscrete<Packet> seaweed0 = Packet.Dry.observation();
			//ObservationDiscrete<Packet> seaweed1 = Packet.Dryish.observation();
			ObservationDiscrete<Packet> seaweed2 = Packet.Damp.observation();
			ObservationDiscrete<Packet> seaweed3 = Packet.Soggy.observation();

			defaultSequences = new ArrayList<ObservationDiscrete<Packet>>();
			defaultSequences.add(seaweed0);
			defaultSequences.add(seaweed2);
			defaultSequences.add(seaweed3);
			defaultHmm = buildHmm();
		}
		return defaultSequences;
	}
	
	static void printObservationSequences(List<ObservationDiscrete<Packet>> sequences) {
		System.out.print("\nObservationSequences: ");
		for (ObservationDiscrete<Packet> sequence : sequences) {
			System.out.print(" >> " + sequence.value.name());
		}
	}
	
	static void runDemo() throws java.io.IOException, CloneNotSupportedException {
		/* Build a HMM and generate observation sequences using this HMM */

		System.out.print("==============================================\n");
		System.out.print(">>>>PresentationDemo\n");
		
		Hmm<ObservationDiscrete<Packet>> hmm = getDefaultHmm();

		List<List<ObservationDiscrete<Packet>>> sequences;
		sequences = generateSequences(hmm);

		/* Baum-Welch learning */

		BaumWelchLearner bwl = new BaumWelchLearner();

		Hmm<ObservationDiscrete<Packet>> learntHmm = buildInitHmm();

		// This object measures the distance between two HMMs
		KullbackLeiblerDistanceCalculator klc = new KullbackLeiblerDistanceCalculator();

		GenericHmmDrawerDot traceAlgo = new GenericHmmDrawerDot();
		
		
		Hmm<ObservationDiscrete<Packet>> preHmm = learntHmm.clone();
		// Incrementally improve the solution
		for (int i = 0; i < 100; i++) {
			/*System.out.println("Distance at iteration " + i + ": "
					+ klc.distance(learntHmm, hmm));*/
			learntHmm = bwl.iterate(learntHmm, sequences);
			
			System.out.println("Distance at iteration " + i + ": "
					+ Math.exp(klc.distance(learntHmm, preHmm)));
			preHmm = learntHmm.clone();
		}

		System.out.println("Resulting HMM:\n" + learntHmm);

		/* Computing the probability of a sequence */

		List<ObservationDiscrete<Packet>> testSequence = getDefaultSequences();
		
		//generate
		//testSequence = testSequence(5);

		System.out.println("Sequence probability: "
				+ learntHmm.probability(testSequence));

		/* Write the final result to a 'dot' (graphviz) file. */

		(new GenericHmmDrawerDot()).write(learntHmm, "learntHmm.dot");
	}
	
	static void runTrainingDemo() throws CloneNotSupportedException {
		/* Build a HMM and generate observation sequences using this HMM */

		System.out.print("\n\n\n\n==============================================\n");
		System.out.print(">>>>PresentationDemo: BaumWelchLearner\n");
		
		Hmm<ObservationDiscrete<Packet>> hmm = getDefaultHmm();

		List<List<ObservationDiscrete<Packet>>> sequences;
		sequences = generateSequences(hmm);

		/* Baum-Welch learning */

		BaumWelchLearner bwl = new BaumWelchLearner();

		Hmm<ObservationDiscrete<Packet>> learntHmm = buildInitHmm();
		System.out.println("Initialize HMM:\n" + learntHmm);

		// This object measures the distance between two HMMs
		KullbackLeiblerDistanceCalculator klc = new KullbackLeiblerDistanceCalculator();

		//GenericHmmDrawerDot traceAlgo = new GenericHmmDrawerDot();
		
		
		Hmm<ObservationDiscrete<Packet>> preHmm = learntHmm.clone();
		// Incrementally improve the solution
		for (int i = 0; i < 20; i++) {
			/*System.out.println("Distance at iteration " + i + ": "
					+ klc.distance(learntHmm, hmm));*/
			learntHmm = bwl.iterate(learntHmm, sequences);
			
			System.out.println("[Kullback Leibler] Distance at iteration " + i + ": "
					+ Math.exp(klc.distance(learntHmm, preHmm)));
			preHmm = learntHmm.clone();
		}

		System.out.println("Resulting HMM:\n" + learntHmm);

		/* Computing the probability of a sequence */

		List<ObservationDiscrete<Packet>> testSequence = getDefaultSequences();
		
		//generate
		//testSequence = testSequence(5);

		// test
		System.out.println("====TEST: ====");
		runViterbiDemo(learntHmm, testSequence, false);
		
		/*System.out.println("Sequence probability: "
				+ learntHmm.probability(testSequence));*/

		/* Write the final result to a 'dot' (graphviz) file. */

		//(new GenericHmmDrawerDot()).write(learntHmm, "learntHmm.dot");
	}
	
	
	static void runForwardDemo() throws IOException {
		System.out.print("\n\n\n\n==============================================\n");
		System.out.print(">>>>runForwardDemo\n");
		Hmm<ObservationDiscrete<Packet>> hmm = getDefaultHmm();
		
		List<ObservationDiscrete<Packet>> testSequence = getDefaultSequences();
		printObservationSequences(getDefaultSequences());
		System.out.println("\nSequence probability: "
				+ hmm.probability(testSequence));
	}
	
	static void runViterbiDemo() {
		System.out.print("\n\n\n\n==============================================\n");
		System.out.print(">>>>runViterbiDemo\n");
		
		System.out.print("...................................\n");
		runViterbiDemo(getDefaultHmm(), getDefaultSequences(), true);
		Random rd = new Random();
		for (int i=0; i<5; i++) {
			System.out.print("...................................\n");
			int n = rd.nextInt(10000) % 3;
			runViterbiDemo(getDefaultHmm(), genObservationSequences(4 + n), false);
		}		
	}
	
	static void runViterbiDemo(
			Hmm<ObservationDiscrete<Packet>> hmm, 
			List<ObservationDiscrete<Packet>> testSequence,
			boolean stepByStep) {
		
		String[] hiddenStateNames = {
			Weather.Sunny.name(),
			Weather.Cloudy.name(),
			Weather.Rainy.name()
		};
		
		printObservationSequences(testSequence);
		
		ViterbiCalculator vcalc = new ViterbiCalculator(testSequence, hmm, stepByStep, hiddenStateNames);
		int[] stateSequences = vcalc.stateSequence();
		System.out.print("\nSate Sequence: ");
		for (int i=0; i<stateSequences.length; i++) {
			System.out.print("  >>  " + Weather.fromInt(stateSequences[i]).name());
		}
		System.out.println(String.format("\nmax Probability: %s", "" + Math.exp(vcalc.lnProbability())));
		System.out.println("Sequence probability: "
				+ hmm.probability(testSequence));
	}

	/* The HMM this example is based on */

	static Hmm<ObservationDiscrete<Packet>> buildHmm() {
		
		int nStates = Weather.values().length;
		int nObservations = Packet.values().length;
		
		double[] initPi = {0.63, 0.17, 0.20};
		double[][] initAij = 
		{ 
				{0.5, 0.25, 0.25},
				{0.375, 0.125, 0.375}, 
				{0.125, 0.675, 0.375} 
		};
		
		double[][] initConfusionMatrix = 
		{
				{0.60, 0.20, 0.15, 0.05},
				{0.25, 0.25, 0.25, 0.25}, 
				{0.05, 0.10, 0.35, 0.50}
		};
		
		return buildHmm(nStates, nObservations, initPi, initAij, initConfusionMatrix);
	}
	
	static Hmm<ObservationDiscrete<Packet>> buildHmm(int nStates, int nObservations, 
			double[] initPi, double[][] initAij, double[][] initConfusionMatrix) {
		
		Hmm<ObservationDiscrete<Packet>> hmm = new Hmm<ObservationDiscrete<Packet>>(
				nStates, new OpdfDiscreteFactory<Packet>(Packet.class));
		
		for (int i=0; i < nStates; i++) {
			hmm.setPi(i, initPi[i]);
		}
		for (int i=0; i < nStates; i++) {
			hmm.setOpdf(
					i,
					new OpdfDiscrete<Packet>(Packet.class, initConfusionMatrix[i]));
		}
		for (int i=0; i < nStates; i++) {
			for (int j=0; j < nStates; j++) {
				hmm.setAij(i, j, initAij[i][j]);
			}
		}

		return hmm;
	}

	/* Initial guess for the Baum-Welch algorithm */
	static double[] genProb(int nItems) {
		double[] vals = new double[nItems];
				
		Random rd = new Random();
		double sum = 0;
		int i=1;
		while (i < nItems) {
			double val = rd.nextDouble();
			if (sum + val >= 1)
				continue;
			
			vals[i] = val;
			sum += val;
			i++;
		}
		vals[0] = 1.0 - sum;
		
		return vals;
	}
	
	static List<ObservationDiscrete<Packet>> genObservationSequences(int nItems) {
		List<ObservationDiscrete<Packet>> sequences = 
			new ArrayList<ObservationDiscrete<Packet>>(nItems);
		
		Random rd = new Random();
		int nObservations = Packet.values().length;
		
		for (int i=0; i<nItems; i++) {
			int val = rd.nextInt(10000) % nObservations;
			sequences.add(Packet.fromInt(val).observation());
		}
		
		return sequences;
	}
	
	static Hmm<ObservationDiscrete<Packet>> buildInitHmm() {
		
		int nStates = Weather.values().length;
		int nObservations = Packet.values().length;
		
		double[] initPi = genProb(nStates);
		double[][] initAij = 
		{ 
				genProb(nStates),
				genProb(nStates), 
				genProb(nStates) 
		};
		double[][] initConfusionMatrix = 
		{
				genProb(nObservations),
				genProb(nObservations),
				genProb(nObservations)
		};
		
		return buildHmm(nStates, nObservations, initPi, initAij, initConfusionMatrix);
	}

	/* Generate several observation sequences using a HMM */

	static <O extends Observation> List<List<O>> generateSequences(Hmm<O> hmm) {
		MarkovGenerator<O> mg = new MarkovGenerator<O>(hmm);

		List<List<O>> sequences = new ArrayList<List<O>>();
		for (int i = 0; i < 200; i++)
			sequences.add(mg.observationSequence(100));

		return sequences;
	}
}
